#include "thread_pool.h"

#include <cassert>

#include <iostream>
#include <sstream>

ThreadPool::ThreadPool(
    size_t thread_n /* = std::thread::hardware_concurrency()*/,
    size_t hold_n /* = 1*/,
    size_t max_n /*= std::thread::hardware_concurrency() * 2*/)
    : shutdown_(true),
      current_thread_count_(thread_n),
      hold_thread_count_(hold_n),
      max_thread_count_(max_n) {}

ThreadPool::~ThreadPool() {
  // ! 这里面不可以调用 Stop()
  // ! 但是内容需要和 Stop() 中一致
  if (!shutdown_) {
    shutdown_.store(true, std::memory_order_relaxed);

    live_map_.clear();

    queue_.Cancel();
  }
}

void ThreadPool::OnMainThreadTick() {
  Slim();
  // TODO 检查是否需要Grow
}

void ThreadPool::Start() {
  assert(shutdown_);
  assert(current_thread_count_ > 0 && hold_thread_count_ > 0 &&
         max_thread_count_ > 0);
  assert(hold_thread_count_ <= current_thread_count_ &&
         hold_thread_count_ <= max_thread_count_ &&
         current_thread_count_ <= max_thread_count_);

  OnStart();

  for (auto i = 0; i < hold_thread_count_; ++i) {
    auto thread = std::thread([this] {
      this->OnThreadStart();

      this->ThreadPool::WorkerLoop();

      this->OnThreadEnd();
    });

    ThreadInfo thread_info(std::move(thread), -1);

    live_map_.emplace(std::make_pair<std::thread::id, ThreadPool::ThreadInfo>(
        thread_info.handler_->get_id(), std::move(thread_info)));
  }

  while (live_map_.size() < current_thread_count_) {
    auto thread = std::thread([this] {
      this->OnThreadStart();

      this->ThreadPool::WorkerLoop();

      this->OnThreadEnd();
    });

    ThreadInfo thread_info(std::move(thread), kThreadPoolWaitMills);

    live_map_.emplace(std::make_pair<std::thread::id, ThreadPool::ThreadInfo>(
        thread_info.handler_->get_id(), std::move(thread_info)));
  }

  shutdown_ = false;
}

void ThreadPool::Stop() {
  assert(!shutdown_);

  OnStop();

  shutdown_.store(true, std::memory_order_relaxed);

  // for (auto &&[_, thread_info] : live_map_) {
  //  if (thread_info.handler_ && thread_info.handler_->joinable()) {
  //    thread_info.handler_->join();
  //  }
  //}

  live_map_.clear();

  queue_.Cancel();
}

void ThreadPool::OnStart() {}

void ThreadPool::OnStop() {}

void ThreadPool::OnThreadStart() {}

void ThreadPool::OnThreadEnd() {}

std::string ThreadPool::PrintState() {
  std::stringstream sstream;

  std::lock_guard<std::mutex> lock(printer_mutex_);

  sstream << "====================\n";

  sstream << "current_thread_count ----- \n";
  sstream << current_thread_count_;
  sstream << "\n";

  sstream << "hold_thread_count ----- \n";
  sstream << hold_thread_count_;
  sstream << "\n";

  sstream << "live map ----- \n";
  for (auto &&[_, thread_info] : live_map_) {
    sstream << "[thread][";
    sstream << thread_info.thread_id_;
    sstream << "] task ";
    sstream << thread_info.task_counter_;
    sstream << "done.";
    sstream << "\n";
  }

  sstream << "====================\n";

  return sstream.str();
}

void ThreadPool::WorkerLoop() {
  while (ProcessNextTask()) {
  }
}

bool ThreadPool::ProcessNextTask() {
  if (auto shutdown = shutdown_.load(std::memory_order_relaxed);
      shutdown && queue_.IsEmpty()) {
    return false;
  }

  auto id = std::this_thread::get_id();

  if (!live_map_.count(id)) {
    return false;
  }

  TaskType task;
  if (queue_.DequeueWaitFor(task, std::chrono::seconds(3))) {
    task();
    {
      std::lock_guard<std::mutex> lock(thread_pool_mutex_);
      if (!live_map_.count(id)) {
        return false;
      }

      auto &thread_info = live_map_[id];
      if (!thread_info.enabled_) {
        return false;
      }
      ++thread_info.task_counter_;
      thread_info.last_busy_time_ = std::chrono::steady_clock::now();
    }

    return true;
  }

  auto now = std::chrono::steady_clock::now();
  {
    std::lock_guard<std::mutex> lock(thread_pool_mutex_);
    if (!live_map_.count(id)) {
      return false;
    }

    auto &thread_info = live_map_[id];
    if (thread_info.wait_mills_ < 0) {
      return true;
    }

    if (!thread_info.enabled_) {
      return false;
    }

    auto diff = std::chrono::duration_cast<std::chrono::milliseconds>(
        now - thread_info.last_busy_time_);
    if (static_cast<int>(diff.count()) >= thread_info.wait_mills_) {
      thread_info.enabled_ = false;
      return false;
    }
  }

  return true;
}

void ThreadPool::Slim() {
  std::lock_guard<std::mutex> lock(thread_pool_mutex_);
  for (auto iter = live_map_.begin(); iter != live_map_.end();) {
    if (iter->second.enabled_) {
      ++iter;
    } else {
      --current_thread_count_;
      iter = live_map_.erase(iter);
    }
  }
}

void ThreadPool::Grow() {}
